import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader
import torchvision
import torchvision.transforms as transforms
from torchvision.models import resnet50
import matplotlib.pyplot as plt
import os
from PIL import Image
import numpy as np
from torch.utils.tensorboard import SummaryWriter
from datetime import datetime

# 定义数据集的配置
class Config:
    IMAGE_SIZE = (224, 224)
    BATCH_SIZE = 32
    NUM_EPOCHS = 25
    LEARNING_RATE = 0.001
    DATA_DIR = "D:/dogs-vs-cats/train"  # 数据集路径
    NUM_CLASSES = 2

# 定义图像预处理
transform_train = transforms.Compose([
    transforms.Resize(Config.IMAGE_SIZE),
    transforms.RandomRotation(15),
    transforms.RandomHorizontalFlip(p=0.5),
    transforms.RandomAffine(
        degrees=0,
        translate=(0.1, 0.1),
        shear=10,
        scale=(0.8, 1.2),
    ),
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])

transform_valid = transforms.Compose([
    transforms.Resize(Config.IMAGE_SIZE),
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])

# 定义数据集类
class CatDogDataset(Dataset):
    def __init__(self, data_dir, transform=None):
        self.data_dir = data_dir
        self.transform = transform
        self.file_paths = []
        self.labels = []
        
        # 获取所有图像文件路径并分类
        for filename in os.listdir(data_dir):
            if filename.startswith('cat'):
                self.file_paths.append(os.path.join(data_dir, filename))
                self.labels.append(0)  # 猫的标签为0
            elif filename.startswith('dog'):
                self.file_paths.append(os.path.join(data_dir, filename))
                self.labels.append(1)  # 狗的标签为1
    
    def __len__(self):
        return len(self.file_paths)
    
    def __getitem__(self, idx):
        image_path = self.file_paths[idx]
        image = Image.open(image_path).convert('RGB')
        label = self.labels[idx]
        
        if self.transform:
            image = self.transform(image)
            
        return image, label

# 创建数据集和数据加载器
train_dataset = CatDogDataset(Config.DATA_DIR, transform=transform_train)
valid_dataset = CatDogDataset(Config.DATA_DIR, transform=transform_valid)

# 使用80%数据用于训练，20%用于验证
train_size = int(0.8 * len(train_dataset))
valid_size = len(train_dataset) - train_size

train_dataset, valid_dataset = torch.utils.data.random_split(train_dataset, [train_size, valid_size])

train_loader = DataLoader(train_dataset, batch_size=Config.BATCH_SIZE, shuffle=True, num_workers=4)
valid_loader = DataLoader(valid_dataset, batch_size=Config.BATCH_SIZE, shuffle=False, num_workers=4)

# 定义模型
class ResNetCatDog(nn.Module):
    def __init__(self, num_classes=Config.NUM_CLASSES):
        super(ResNetCatDog, self).__init__()
        self.resnet = resnet50(pretrained=True)
        self.resnet.fc = nn.Linear(self.resnet.fc.in_features, num_classes)
        
    def forward(self, x):
        return self.resnet(x)

model = ResNetCatDog()

# 强制使用 CPU
device = torch.device("cpu")
model.to(device)

# 定义损失函数和优化器
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=Config.LEARNING_RATE)

# 使用TensorBoard记录训练过程
current_time = datetime.now().strftime("%Y%m%d-%H%M%S")
log_dir = os.path.join("logs", current_time)
writer = SummaryWriter(log_dir)

# 训练函数
def train_model(model, criterion, optimizer, num_epochs=Config.NUM_EPOCHS):
    train_losses = []
    train_accuracies = []
    valid_losses = []
    valid_accuracies = []
    
    for epoch in range(num_epochs):
        # 训练阶段
        model.train()
        running_loss = 0.0
        running_corrects = 0
        
        for inputs, labels in train_loader:
            inputs = inputs.to(device)
            labels = labels.to(device)
            
            optimizer.zero_grad()
            
            outputs = model(inputs)
            loss = criterion(outputs, labels)
            
            loss.backward()
            optimizer.step()
            
            _, preds = torch.max(outputs, 1)
            running_loss += loss.item() * inputs.size(0)
            running_corrects += torch.sum(preds == labels.data)
        
        epoch_loss = running_loss / len(train_loader.dataset)
        epoch_acc = running_corrects.double() / len(train_loader.dataset)
        
        train_losses.append(epoch_loss)
        train_accuracies.append(epoch_acc)
        
        # 验证阶段
        model.eval()
        running_loss = 0.0
        running_corrects = 0
        
        with torch.no_grad():
            for inputs, labels in valid_loader:
                inputs = inputs.to(device)
                labels = labels.to(device)
                
                outputs = model(inputs)
                loss = criterion(outputs, labels)
                
                _, preds = torch.max(outputs, 1)
                running_loss += loss.item() * inputs.size(0)
                running_corrects += torch.sum(preds == labels.data)
        
        epoch_loss = running_loss / len(valid_loader.dataset)
        epoch_acc = running_corrects.double() / len(valid_loader.dataset)
        
        valid_losses.append(epoch_loss)
        valid_accuracies.append(epoch_acc)
        
        print(f'Epoch {epoch+1}/{num_epochs}')
        print(f'Train Loss: {epoch_loss:.4f}, Train Acc: {epoch_acc:.4f}')
        print(f'Valid Loss: {epoch_loss:.4f}, Valid Acc: {epoch_acc:.4f}')
        print('---' * 10)
        
        # 记录到TensorBoard
        writer.add_scalar('Loss/train', epoch_loss, epoch)
        writer.add_scalar('Accuracy/train', epoch_acc, epoch)
        writer.add_scalar('Loss/valid', epoch_loss, epoch)
        writer.add_scalar('Accuracy/valid', epoch_acc, epoch)
    
    writer.close()
    return train_losses, train_accuracies, valid_losses, valid_accuracies

# 可视化函数
def visualize_results(train_losses, train_accuracies, valid_losses, valid_accuracies):
    plt.figure(figsize=(10, 5))
    plt.subplot(1, 2, 1)
    plt.plot(train_losses, label='Train Loss')
    plt.plot(valid_losses, label='Valid Loss')
    plt.title('Loss Curves')
    plt.xlabel('Epochs')
    plt.ylabel('Loss')
    plt.legend()
    
    plt.subplot(1, 2, 2)
    plt.plot(train_accuracies, label='Train Accuracy')
    plt.plot(valid_accuracies, label='Valid Accuracy')
    plt.title('Accuracy Curves')
    plt.xlabel('Epochs')
    plt.ylabel('Accuracy')
    plt.legend()
    
    plt.tight_layout()
    plt.show()

# 训练模型
train_losses, train_accuracies, valid_losses, valid_accuracies = train_model(model, criterion, optimizer)

# 可视化训练结果
visualize_results(train_losses, train_accuracies, valid_losses, valid_accuracies)